import sys
sys.path.append('/home/A03_tmpdata0/voice-gender-classifier-main')
import os
import torch
import multiprocessing
from multiprocessing import Pool, Lock
from model import ECAPA_gender
from gxl_ai_utils.utils import utils_file

# 真实文件在：/home/A03_tmpdata0/voice-gender-classifier-main
# ========= 配置 =========
model_path = "/home/A03_tmpdata0/voice-gender-classifier-main/JaesungHuh/voice-gender-classifier"

input_data_list_path = "/home/A02_tmpdata3/osum_s2s/S2SChat_osum_setting_qa_527_updated_by_cywang_added_by_20250616/raw_data/data.list"
input_data_list_path_with_gender = "/home/A02_tmpdata3/osum_s2s/S2SChat_osum_setting_qa_527_updated_by_cywang_added_by_20250616/raw_data/data_with_gender.list"
dict_list = utils_file.load_dict_list_from_jsonl(input_data_list_path)
wav_list = [d["wav"] for d in dict_list]
output_wav_path = "/home/A02_tmpdata3/tmp.list"
utils_file.write_list_to_file(wav_list, output_wav_path)
path_file = output_wav_path
output_file = "/home/A02_tmpdata3/osum_s2s/S2SChat_osum_setting_qa_527_updated_by_cywang_added_by_20250616/raw_data/gender.list"
num_processes = 80  # 总进程数
num_gpus = 8  # 8张GPU卡

# ========= 共享对象 =========
lock = Lock()

# ========= 子进程初始化（通过进程ID绑定GPU） =========
def init_worker():
    # 获取当前进程ID（在进程池内的索引）
    global model, device
    proc_id = multiprocessing.current_process()._identity[0] - 1  # 进程池内ID从0开始
    gpu_id = proc_id % num_gpus  # 每个进程绑定一个GPU（0-7循环）
    device = torch.device(f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu")
    # 加载模型到当前GPU
    model_local = ECAPA_gender.from_pretrained(model_path)
    model_local.eval()
    model_local.to(device)
    model = model_local

# ========= 推理函数 =========
def process_wav(wav_path):
    global model, device
    uid = os.path.splitext(os.path.basename(wav_path))[0]
    try:
        with torch.no_grad():
            gender = model.predict(wav_path, device=device)
        line = f"{uid} {wav_path} {gender}\n"
        # 加锁写入避免冲突
        with lock:
            with open(output_file, "a", encoding="utf-8") as f:
                f.write(line)
    except Exception as e:
        print(f"Error processing {wav_path}: {e}")

# ========= 主程序 =========
if __name__ == "__main__":
    # 读取wav路径
    with open(path_file, "r", encoding="utf-8") as f:
        lines = [line.strip() for line in f if line.strip()]
        wav_paths = [line.split()[0] for line in lines]

    # 过滤已处理的UID
    completed_uids = set()
    if os.path.exists(output_file):
        with open(output_file, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    completed_uids.add(line.split()[0])
    remaining_wavs = [p for p in wav_paths if os.path.splitext(os.path.basename(p))[0] not in completed_uids]

    print(f"Total: {len(wav_paths)}, Completed: {len(completed_uids)}, Remaining: {len(remaining_wavs)}")

    if remaining_wavs:
        # 进程池初始化时不传递参数，由进程自身获取ID绑定GPU
        with Pool(processes=num_processes, initializer=init_worker) as pool:
            pool.map(process_wav, remaining_wavs)
    else:
        print("All files processed.")
    print(f"Gender inference results saved to {output_file}")
    gender_dict = utils_file.load_dict_from_scp(output_file)
    new_gender_dict = {}
    for key in gender_dict:
        new_gender_dict[key] = "<" + gender_dict[key].strip().split()[1].upper() + ">"
    new_dict_list = []
    for dict_i in dict_list:
        if dict_i['key'] in new_gender_dict:
            dict_i['extra']['gender'] = new_gender_dict[dict_i['key']]
        else:
            continue
        new_dict_list.append(dict_i)
    utils_file.write_dict_list_to_jsonl(new_dict_list, input_data_list_path_with_gender)

